from datetime import datetime, timedelta
from ccxt.base.exchange import Exchange
from backtrader.utils.datehelper import str2date, date2str, str2timestamp, timestamp2str


class TradingCalendar:
    def __init__(self, start_date_str: str, end_date_str: str, time_frame='1d', offset=2):
        self.format_str = "%Y-%m-%d %H:%M:%S"
        self.timeframe = time_frame.lower()
        self.start_date = str2date(start_date_str, self.format_str)
        _end_date = str2date(end_date_str, self.format_str)
        self.duration = Exchange.parse_timeframe(self.timeframe)
        roll = int((_end_date.timestamp() - self.start_date.timestamp()) / self.duration)
        self.end_date = self.start_date + timedelta(seconds=self.duration * (roll - offset))

    def next(self, date_str, roll=1):
        _date = str2date(date_str, self.format_str)
        if not self.is_fit(_date.timestamp()):
            raise Exception(f"{date_str} is not trading Calendar point.")

        if roll > 0:
            next_date = _date + timedelta(seconds=self.duration * roll)
        else:
            next_date = _date
        return next_date.strftime(self.format_str)

    def first(self):
        return self.start_date.strftime(self.format_str)

    def last(self):
        return self.end_date.strftime(self.format_str)

    def latest(self):
        """
        最新的日历
        """
        roll = int((datetime.today().timestamp() - self.start_date.timestamp()) / self.duration)
        latest_calendar = self.start_date + timedelta(seconds=self.duration * roll)
        return latest_calendar.strftime(self.format_str)

    def is_fit(self, timestamp):
        return (timestamp - self.start_date.timestamp()) % self.duration == 0

    def adjust(self, timestamp):
        roll = int((timestamp - self.start_date.timestamp()) / self.duration)
        return self.start_date.timestamp() + self.duration * roll

    def head(self):
        return self.range_list(date2str(self.start_date), None, limit=10)

    def tail(self):
        return self.range_list(None, date2str(self.end_date), limit=10)

    def len(self, start, end):
        _start = str2timestamp(start)
        _end = str2timestamp(end)

        if not self.is_fit(_start):
            raise Exception(f"{start} is not trading Calendar point.")

        if not self.is_fit(_end):
            raise Exception(f"{end} is not trading Calendar point.")

        return int((_end - _start) / self.duration) + 1

    def range_list(self, start, end, limit=None):
        """
        start和end都包含
        """
        _start = str2timestamp(start) if start else None
        _end = str2timestamp(end) if end else None

        if _start is None and _end is None:
            raise Exception(f" start and end must be not none.")

        if (_start is None or _end is None) and limit is None:
            raise Exception(f" start and end have none, limit must not be none.")

        if _start and _end is None:
            _end = _start + self.duration * (limit - 1)

        if _start is None and _end:
            _start = _end - self.duration * (limit - 1)

        if not self.is_fit(_start):
            raise Exception(f"{start} is not trading Calendar point.")

        if _end and not self.is_fit(_end):
            raise Exception(f"{end} is not trading Calendar point.")

        if _end and _start > _end:
            raise Exception(f"start {start} must be less than end {end}.")

        rs = []
        roll = 0
        while True:
            _d = _start + self.duration * roll
            if limit and roll > limit:
                break
            if _end and _d > _end:
                break
            rs.append(_d)
            roll += 1

        return list(map(lambda x: timestamp2str(x), rs)) if rs else []
